import os
import json
import runpy
import traceback
from pathlib import Path

g_tool_set = {}
g_tool_functions = {}

def register_tool(tool):
    def decorator(function):
        name = function.__name__
        g_tool_set[name] = tool
        g_tool_functions[name] = function
        return function
    return decorator

def check_tool_name(name):
    if '/' in name or '\\' in name:
        raise RuntimeError(f'invalid tool name {name}')

category_list = ["web", "clipboard", "file", "media", "coding", "shell", "image", "IoT", "assistant", "tooling", "other"]

# @register_tool({
#     "description": "Load tools by category",
#     "parameters": {
#         "type": "object",
#         "properties": {
#             "category": {
#                 "type": "string",
#                 "enum": category_list
#             }
#         },
#         "required": [
#             "category"
#         ]
#     },
#     "category": "tooling"
# })
# def load_tools(category):
#     g_category_enabled.add(category)
#     tools = []
#     for name, descriptor in g_tool_set.items():
#         if descriptor.get("category", "other") == category:
#             tools.append(name)
#     return {"success": True, "tools": tools}

@register_tool({
    "description": "Create a tool for future use",
    "parameters": {
        "type": "object",
        "properties": {
            "name": {
                "type": "string",
                "description": "name of the tool, e.g.: read_file"
            },
            "description": {
                "type": "string",
                "description": "description of what the tool do, e.g.: read the content of a given file path"
            },
            "code": {
                "type": "string",
                "description": "a python function to handle tool request, e.g.:\nimport pathlib\n\ndef read_file(file_path: str):\n    path_obj = pathlib.Path(file_path)\n    if path_obj.exists() and path_obj.is_file():\n        return {'success': False, 'message': 'File does not exist or is not a file.'}\n    with open(file_path, 'r') as file:\n        content = file.read()\n    return {'success': True, 'content': content}"
            },
            "parameters": {
                "type": "string",
                "description": "describe the parameters of the function in the JSON Schema format, e.g.:\n{\"file_path\": {\"type\": \"string\"}}"
            },
            "category": {
                "type": "string",
                "enum": category_list
            }
        },
        "required": [
            "name",
            "descrption",
            "code",
            "parameters",
            "category"
        ]
    },
    "category": "tooling"
})
def create_tool(name, description, code, parameters, category="other"):
    check_tool_name(name)
    tools_dir = get_tools_dir()
    if name in g_tool_functions:
        del g_tool_functions[name]
    with open(tools_dir / (name + '.py'), 'w') as f:
        f.write(code)
    load_tool_function(name)
    parameters = json.loads(parameters)
    assert isinstance(parameters, dict)
    required = list(parameters.keys())
    descriptor = {
        "description": description,
        "parameters": {
            "type": "object",
            "properties": parameters,
            "required": required,
        },
        "category": category
    }
    with open(tools_dir / (name + '.json'), 'w') as f:
        json.dump(descriptor, f, ensure_ascii=False)
    g_tool_set[name] = descriptor
    return {"success": True}

@register_tool({
    "description": "Delete a tool",
    "parameters": {
        "type": "object",
        "properties": {
            "name": {
                "type": "string",
                "description": "name of the tool"
            }
        },
        "required": [
            "name"
        ]
    },
    "category": "tooling"
})
def delete_tool(name):
    check_tool_name(name)
    tools_dir = get_tools_dir()
    if (tools_dir / (name + '.json')).exists():
        os.unlink(tools_dir / (name + '.json'))
        os.unlink(tools_dir / (name + '.py'))
        g_tool_functions.pop(name, None)
        return {'success': True}
    return {'success': False, 'message': 'tool not exist or not deletable'}

@register_tool({
    "description": "List available tools",
    "parameters": {
        "type": "object",
        "properties": {
        }
    },
    "category": "tooling"
})
def list_tools():
    return g_tool_set

@register_tool({
    "description": "Get the source code of a tool",
    "parameters": {
        "type": "object",
        "properties": {
            "name": {
                "type": "string",
                "description": "name of the tool"
            }
        },
        "required": [
            "name"
        ]
    },
    "category": "tooling"
})
def get_tool_source(name):
    check_tool_name(name)
    tools_dir = get_tools_dir()
    if (tools_dir / (name + '.json')).exists():
        with open(tools_dir / (name + '.py')) as f:
            content = f.read()
        return {'success': True, 'source': content}
    return {'success': False, 'message': 'tool not exist or builtin'}

def get_tools_dir():
    from . import __file__ as root_mod_path
    return Path(os.path.dirname(root_mod_path)) / 'tools'

def load_tool_function(name):
    if name in g_tool_functions:
        function = g_tool_functions[name]
    else:
        tools_dir = get_tools_dir()
        module = runpy.run_path(tools_dir / (name + '.py'))
        assert name in module, f'symbol {name} not found in the tool code'
        function = module[name]
        g_tool_functions[name] = function
    return function

def load_all_tools():
    tools_dir = get_tools_dir()
    if not tools_dir.exists():
        os.mkdir(tools_dir)
    for e in os.scandir(tools_dir):
        if e.name.endswith('.json'):
            try:
                with open(e.path, 'r') as f:
                    descriptor = json.load(f)
                name = e.name[:-len('.json')]
                check_tool_name(name)
                load_tool_function(name)
                g_tool_set[name] = descriptor
            except:
                print(f'Error loading tool {e.name}')
                traceback.print_exc()
                continue

def list_tool_descriptors(categories=None):
    # for name, descriptor in g_tool_set.items():
    #     if not 'category' in descriptor:
    #         raise RuntimeError(name)
    result = []
    for name, descriptor in g_tool_set.items():
        if categories is not None and descriptor.get("category", "other") not in categories:
            continue
        function = {"name": name}
        if "descrption" in descriptor:
            function["descrption"] = descriptor["descrption"]
        if "parameters" in descriptor:
            function["parameters"] = descriptor["parameters"]
        result.append({"type": "function", "function": function})
    if not result:
        return None
    return result

def get_tool_function(name):
    return g_tool_functions.get(name, None)

# {"file_path": {"type": "string"}, "temperature": {"type": "integer"}, "unit": {"enum": ["celsius", "fahrenheit"]}, "cities": {"type": "array", "items": {"type": "string"}}, "percentage": {"type": "number", "minimum": 0, "maximum": 100}}
# "import requests\nimport json\nimport urllib.request\n\n# Function to download Bilibili video by bvid\n# Reference: https://github.com/SocialSisterYi/bilibili-API-collect\n\ndef download_bilibili_video(bvid: str, save_path: str):\n try:\n api_url = f'https://api.bilibili.com/x/web-interface/view?bvid={bvid}'\n response = requests.get(api_url)\n data = response.json()\n video_title = data['data']['title']\n video_url = data['data']['pages'][0]['part']\n video_download_url = data['data']['pages'][0]['video_download_url']\n\n file_path = f'{save_path}/{video_title}.mp4'\n urllib.request.urlretrieve(video_download_url, file_path)\n return {'success': True, 'message': f'Video \"{video_title}\" downloaded successfully.'}\n except Exception as e:\n return {'success': False, 'message': str(e)}"}

load_all_tools()
